Conversation
- Rename src/loremjax/ to src/lorem/ so `import lorem` works - Move standalone train.py into package as lorem.train with `lorem-train` console script - Add train-bec example with prepare.py, model/settings YAML, and run.sh - Factor shared model code (Initial, MLP, Update, RadialCoefficients, spherical helpers) into models/backbone.py to deduplicate mlip.py and bec.py - Fix __init__.py exports (LOREMCalculator, LoremBEC, comms) - Add missing dependencies to pyproject.toml (e3x, flax, optax, jaxtyping) - Adopt ruff rules from marathon (relaxed E741/E731/E402/E501 ignores) - Update test to use new package name Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Small non-periodic C9H4 dataset (21 structures), trains Lorem model for 2 epochs on energy and forces. Runs in ~45 seconds. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
|
Added This comment was written by Claude Code on behalf of @sirmarcel. |
- Make backbone importable without marathon (lazy import of `masked`) - Skip test_init when comms is not installed (CI) - Fix stale loremjax references in _install_ipi_driver.py and test_ipi.py - Add test_backbone.py: tests for spherical helpers (degree_wise_trace, degree_wise_repeat, spherical_norm + custom JVP gradient check), MLP, ChemicalEmbedding, RadialEmbedding, and Initial module Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
|
Note: This PR cannot be fully merged until
This comment was written by Claude Code on behalf of @sirmarcel. |
jaxpme is public so we can pull it from GitHub directly. The comms and marathon imports are wrapped in try/except so that backbone tests can run in CI without them. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
|
Fix CI failures (55f1c0d, written by Claude) The backbone tests were failing in CI because Changes:
Remaining CI blocker: |
pip rejects 'jaxpme' because the package metadata declares 'jax-pme'. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
- models/__init__.py: wrap bec/mlip imports in try/except - mlip.py, bec.py: use _masked from backbone instead of direct marathon.utils import Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
The calculator had broken imports (nonexistent SetUpEwald, get_batch) and wrong field access patterns. Rewired to use lorem.batching directly, added from_model() classmethod for easy construction without checkpoints. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Also fixes the i-PI driver command name (loremjax-install-ipi-driver -> lorem-install-ipi-driver) to match pyproject.toml. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Previously these were hardcoded to None and always auto-detected from folder structure. Now settings.yaml can override them while keeping the folder-based fallback as default. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Wire up the Calculator → i-PI BEC data path: extract "apt" from LoremBEC model output, reshape to (3*natoms, 3), and expose as "BEC" for i-PI's driven_dynamics module. Implement LOREMCalculator as a proper Calculator subclass and clean up the i-PI driver (remove buggy compute_structure override that operated on JSON strings). Add md-ipi example with input.xml (eda-nve, bec mode="driver"), run.sh, start.xyz, and README. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
04291dd to
7b5b3f6
Compare
7b5b3f6 to
1a57216
Compare
GetPropertiesMixin lacks calculate_properties(), which i-PI's ASEDriver requires. BaseCalculator provides it (and inherits from GetPropertiesMixin), so this is a strict superset. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
The model predicts the virial (stress * volume) as a 3x3 tensor. Convert to ASE convention: divide by cell volume and return as Voigt 6-element array. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
jax-pme moved halfspace handling into prepare(), so get_batch no longer accepts this parameter. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
- Update BEC tests to use "born_effective_charges" key and (N,3,3) shape to match the ASE-compatible calculator output - Export LOREMCalculator alias from __init__.py - (ipi entry point was fixed by reinstall after prior flattening commit) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
hermes is deprecated; marathon.grain is now the canonical module. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Add marathon (release/v0.2.0) as a proper git dependency instead of requiring separate local install. Fix ToBatch.k_size to use kg.shape[0] directly instead of count_halfspace_kvectors(kg.shape), which expects grid dimensions (nx, ny, nz) but receives the flat k-vector array shape. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
- Add marathon[grain] extra to pyproject.toml - Replace example run.sh scripts with README docs - Add tox -e examples environment for smoke-testing examples - Add examples CI job to tests workflow - Remove stale marathon-dev install note from README - Add tox commands to README development section - Fix isort in ipi.py (drop unneeded __future__ import) - Add uv.lock to .gitignore Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
uvloop (transitive dep via flax/orbax) does not support Windows. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
- Document full training workflow (data prep, config, running) - Explain $DATASETS environment variable - Add collapsible reference table for all training settings - Default matmul precision to float32 instead of default Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Both are hard dependencies — the guards were dead code. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Summary
loremjax→loremsoimport loremworks as expectedtrain.pyinto the package aslorem.trainwith alorem-trainconsole script entry pointtrain-becexample (BEC prediction training)models/backbone.py, deduplicatingmlip.pyandbec.pytox -e examplessmoke test environmentipi.py, adduv.lockto.gitignoreTest plan
tox -e lintpassestox -e testspasses (18 tests)tox -e examplespasses (calculator, train-mlp, train-bec)🤖 Generated with Claude Code